package com.langll.ai.controller;

import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.langll.ai.chat.ChatSession;
import com.langll.ai.chat.ChatSessionService;
import com.langll.ai.chat.RedisChatMemory;
import com.langll.ai.controller.request.DeleteSessionRequest;
import com.langll.ai.controller.request.GenerateRequest;
import com.langll.ai.controller.response.ModelTagsResponse;
import com.langll.ai.controller.response.GenerateStreamResponse;
import com.langll.ai.service.ModelService;
import io.swagger.v3.oas.annotations.OpenAPIDefinition;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.info.Info;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.UUID;
import java.time.Instant;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @author langll
 * @des
 * @date 2025/2/6 上午10:47
 */
@RestController
@RequestMapping("/chat")
@OpenAPIDefinition(info = @Info(title = "Chat API", version = "1.0", description = "API for chat operations"))
@Tag(name = "AI聊天", description = "集成机器人角色设定，会话存储，聊天记录持久化")
@RequiredArgsConstructor
public class AIChatController {

    private final ChatClient chatClient;
    private final RedisChatMemory redisChatMemory;
    private final ChatSessionService chatSessionService;
    private final ModelService modelService;

    @Operation(summary = "流式回答聊天")
    @GetMapping(value = "/ai/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatResponse> generateStream(
            @RequestParam(value = "message", defaultValue = "portfolio") String message, @RequestParam String sessionId,
            @RequestParam String userId) {
        Assert.notNull(message, "message不能为空");
        Assert.notNull(userId, "userId不能为空");
        // 默认生成一个会话
        if (StrUtil.isBlank(sessionId)) {
            sessionId = UUID.randomUUID().toString();
            ChatSession chatSession = new ChatSession().setSessionId(sessionId)
                    .setSessionName(message.length() >= 15 ? message.substring(0, 15) : message);
            chatSessionService.saveSession(chatSession, userId);
        }
        String finalSessionId = sessionId;
        return chatClient.prompt().user(message).advisors(advisorSpec -> advisorSpec
                .param(CHAT_MEMORY_CONVERSATION_ID_KEY, finalSessionId).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .stream().chatResponse();
    }

    @Operation(summary = "获取聊天记录")
    @GetMapping("/ai/messages")
    public ResponseEntity<List<Message>> getMessages(@RequestParam String sessionId) {
        Assert.notNull(sessionId, "sessionId不能为空");
        return ResponseEntity.ok(redisChatMemory.get(sessionId, 10));
    }

    @Operation(summary = "获取会话列表")
    @GetMapping("/ai/sessions")
    public ResponseEntity<List<ChatSession>> getSessions(@RequestParam String userId) {
        Assert.notNull(userId, "userId不能为空");
        return ResponseEntity.ok(chatSessionService.getSessions(userId));
    }

    @Operation(summary = "删除会话")
    @PostMapping("/ai/deleteSession")
    public ResponseEntity<Boolean> deleteSession(@RequestBody DeleteSessionRequest request) {
        Assert.notNull(request.getUserId(), "userId不能为空");
        Assert.notNull(request.getSessionId(), "sessionId不能为空");
        boolean result = chatSessionService.deleteSession(request.getUserId(), request.getSessionId());
        return ResponseEntity.ok(result);
    }

    @Operation(summary = "普通聊天")
    @GetMapping(value = "/ai/generate")
    public ResponseEntity<String> generate(@RequestParam(value = "message", defaultValue = "讲个笑话") String message,
            @RequestParam String sessionId, @RequestParam String userId) {
        Assert.notNull(message, "message不能为空");
        Assert.notNull(userId, "userId不能为空");
        // 默认生成一个会话
        if (StrUtil.isBlank(sessionId)) {
            sessionId = UUID.randomUUID().toString();
            ChatSession chatSession = new ChatSession().setSessionId(sessionId)
                    .setSessionName(message.length() >= 15 ? message.substring(0, 15) : message);
            chatSessionService.saveSession(chatSession, userId);
        }
        String finalSessionId = sessionId;
        return ResponseEntity.ok(chatClient.prompt().user(message).advisors(advisorSpec -> advisorSpec
                .param(CHAT_MEMORY_CONVERSATION_ID_KEY, finalSessionId).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .call().content());
    }

    @Operation(summary = "获取模型信息")
    @GetMapping("/ai/tags")
    public ResponseEntity<ModelTagsResponse> getModelTags() {
        return ResponseEntity.ok(modelService.getModelTags());
    }

    @Operation(summary = "AI生成回答（POST请求，支持流式输出）")
    @PostMapping(value = "/ai/generate", consumes = MediaType.APPLICATION_JSON_VALUE, produces = MediaType.APPLICATION_OCTET_STREAM_VALUE)
    public Flux<String> generateWithPost(@RequestBody GenerateRequest request) {
        Assert.notNull(request, "请求参数不能为空");
        Assert.notNull(request.getPrompt(), "prompt不能为空");

        // 检查是否启用流式输出
        boolean enableStream = request.getStream() != null && request.getStream();
        String sessionId = request.getSessionId();
        String userId = request.getUserId();
        if (StrUtil.isBlank(userId)) {
            userId = "defaultUser"; // 默认用户ID
        }

        String message = request.getPrompt();
        String model = StrUtil.isNotBlank(request.getModel()) ? request.getModel() : "qwen3:0.6b";

        if (StrUtil.isBlank(sessionId)) {
            sessionId = UUID.randomUUID().toString();
        }

        String finalSessionId = sessionId;
        // 默认生成一个会话
        if (!chatSessionService.getSessions(userId).stream()
                .anyMatch(session -> session.getSessionId().equals(finalSessionId))) {

            ChatSession chatSession = new ChatSession().setSessionId(finalSessionId)
                    .setSessionName(message.length() >= 15 ? message.substring(0, 15) : message);
            chatSessionService.saveSession(chatSession, userId);
        }

        if (enableStream) {
            // 流式输出 - 每个数据块都是一行JSON
            return chatClient.prompt()
                    .user(message)
                    .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, finalSessionId)
                            .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                    .stream()
                    .content()
                    .map(content -> {
                        GenerateStreamResponse response = new GenerateStreamResponse()
                                .setCreatedAt(Instant.now().toString())
                                .setDone(false)
                                .setModel(model)
                                .setResponse(content);
                        return convertToJsonLine(response);
                    })
                    .concatWith(Flux.just(convertToJsonLine(new GenerateStreamResponse()
                            .setCreatedAt(Instant.now().toString())
                            .setDone(true)
                            .setModel(model)
                            .setResponse(""))));
        } else {
            // 非流式输出
            String response = chatClient.prompt()
                    .user(message)
                    .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, finalSessionId)
                            .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                    .call()
                    .content();
            return Flux.just(convertToJsonLine(new GenerateStreamResponse()
                    .setCreatedAt(Instant.now().toString())
                    .setDone(true)
                    .setModel(model)
                    .setResponse(response)));
        }
    }

    private String convertToJsonLine(GenerateStreamResponse response) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            // 将响应对象转换为JSON字符串并添加换行符
            String jsonString = objectMapper.writeValueAsString(response);
            // 返回JSON字符串 + 换行符，方便前端按行解析
            return jsonString + "\n";
        } catch (JsonProcessingException e) {
            e.printStackTrace();
            return "{\"error\":\"JSON conversion failed\"}\n";
        }
    }

}